import tensor as ts
import function as func


tensors = []
variable = []

def addToTensors(*tList):
	global tensors
	for tensor in tList:
		tensors.append(tensor)

def addToVariable(*tList):
	global variable
	for tensor in tList:
		variable.append(tensor)


class linear:
	def __init__(self,inputSize,outputSize,activation_func=None):
		self.W = ts.Tensor((inputSize,outputSize))
		self.B = ts.Tensor((1,outputSize))
		self.W.random_normal(0,1)
		self.B.zeros()
		if activation_func:
			self.activation_func = func.funcDict[activation_func]
		addToVariable(self.W,self.B)

	def __call__(self,inputTensor):
		a = inputTensor.matmul(self.W)
		b = a+self.B
		# c = self.activation_func(b)
		try:
			c = self.activation_func(b)
			addToTensors(a,b,c,self.W,self.B)
			return c
		except AttributeError:
			addToTensors(a,b,self.W,self.B)
			return b
		
class MSELoss:
	def __call__(self,out,feature):
		loss = out.MSELoss(feature)
		addToTensors(loss)
		return loss

class softmax_crossEntropy:
	def __call__(self,out,feature):
		loss = out.softmax_crossEntropy(feature)
		addToTensors(loss)
		return loss






























































